#include <boost/graph/adjacency_list.hpp>
#include <boost/graph/kruskal_min_spanning_tree.hpp>
#include <iostream>
#include <vector>
#include <random>
#include <fstream>

using namespace boost;

typedef adjacency_list <vecS, vecS, undirectedS, no_property,
property<edge_weight_t, int>>
        Graph;
typedef graph_traits<Graph>::edge_descriptor Edge;
typedef std::pair<int, int> E;

std::vector<std::pair<int, int>> get_edge_array(std::vector<std::vector<int>> &graph) {
    std::vector<std::pair<int, int>> edge_array;
    int n_vertex = graph.size();
    for (int i = 0; i < n_vertex; i++) {
        int n_neighbor = graph[i].size();
        for (int j = 0; j < n_neighbor; j++) {
            if (i < graph[i][j]) {
                edge_array.push_back(std::make_pair(i, graph[i][j]));
            }
        }
    }
    return edge_array;
}

std::vector<int> get_weight(int n_edge, int random_seed) {
    std::vector<int> weight_l(n_edge);
    srand(random_seed);
    for (int i = 0; i < n_edge; i++) {
        weight_l[i] = rand() % 100;
    }
    return weight_l;
}

std::vector<std::vector<int>> get_span_tree(std::vector<std::vector<int>> &graph, int random_seed) {
    std::vector<std::pair<int, int>> edge_array = get_edge_array(graph);
    int n_edge = edge_array.size();
    int n_vertex = graph.size();
    std::vector<int> weight_l = get_weight(n_edge, random_seed);

    Graph g(edge_array.data(), edge_array.data() + n_edge, weight_l.data(), n_vertex);
    property_map<Graph, edge_weight_t>::type weight = get(edge_weight, g);
    std::vector<Edge> tmp_span_tree;

    kruskal_minimum_spanning_tree(g, std::back_inserter(tmp_span_tree));

    std::vector<std::vector<int>> span_tree(n_vertex);

    std::cout << "Print the edges in the MST:" << std::endl;
    for (std::vector<Edge>::iterator iter = tmp_span_tree.begin();
         iter != tmp_span_tree.end(); ++iter) {
        std::cout << source(*iter, g) << " <--> " << target(*iter, g)
                  << " with weight of " << weight[*iter] << std::endl;
        span_tree[source(*iter, g)].push_back(target(*iter, g));
        span_tree[target(*iter, g)].push_back(source(*iter, g));
    }
    return span_tree;
}

int main() {
    using namespace std;

    vector<vector<int>> graph(5);
    graph[0].push_back(2);
    graph[0].push_back(4);
    graph[1].push_back(3);
    graph[1].push_back(4);
    graph[1].push_back(2);
    graph[2].push_back(0);
    graph[2].push_back(1);
    graph[2].push_back(3);
    graph[3].push_back(4);
    graph[3].push_back(1);
    graph[4].push_back(1);
    graph[4].push_back(3);
    graph[4].push_back(0);
    std::vector<std::vector<int>> span_tree = get_span_tree(graph, 2);
    for (int i = 0; i < 5; i++) {
        int n_neighbor = span_tree[i].size();
        for (int j = 0; j < n_neighbor; j++) {
            printf("%d ", span_tree[i][j]);
        }
        printf("\n");
    }
    return 0;
}